from openai import OpenAI
from zhipuai import ZhipuAI
import os
import requests
from bs4 import BeautifulSoup
import re
import json
from pydantic import BaseModel

class NewsItem(BaseModel):
    """
    新闻条目
    包含新闻标题和链接

    参数:
        title: 新闻标题
        url: 新闻链接
    """
    title: str
    url: str

def call_llm(messages):
    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "sk-8232f1f94ce448f0b5d71563ee23969b"),base_url="https://api.deepseek.com/v1")
    
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=messages,
        temperature=0.7
    )
    
    return response.choices[0].message.content

def call_zhipuai(prompt_or_messages) -> str:
    """
    调用智谱AI的API

    参数:
        prompt_or_messages: 可以是字符串提示或对话历史列表

    返回:
        response.choices[0].message.content: 智谱AI的回答
    """
    import time
    
    # 获取API密钥
    api_key = os.environ.get("ZHIPUAI_API_KEY", "9e2458a8acf46274d1c1b9418eec500c.oBojpibVHt3IyXnz")
    client = ZhipuAI(api_key=api_key)
    
    # 检查输入类型并转换为正确的消息格式
    if isinstance(prompt_or_messages, str):
        messages = [{"role": "user", "content": prompt_or_messages}]
    else:
        # 确保messages是正确的列表格式
        if not isinstance(prompt_or_messages, list):
            print(f"警告: 输入格式错误，类型为 {type(prompt_or_messages)}")
            messages = [{"role": "user", "content": str(prompt_or_messages)}]
        else:
            messages = prompt_or_messages
    
    # 确保每条消息都有正确的格式
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
            print(f"警告: 消息格式错误: {msg}")
            messages[i] = {"role": "user", "content": str(msg)}
    
    # 最大重试次数
    max_retries = 3
    base_wait_time = 2  # 基础等待时间，秒
    
    for attempt in range(max_retries):
        try:
            print(f"尝试调用智谱AI，第 {attempt+1} 次...")
            
            # 打印消息内容以便调试
            if attempt > 0:  # 只有在重试时才打印
                print(f"消息内容: {messages}")
            
            response = client.chat.completions.create(
                model="glm-4-flash",
                messages=messages,
                temperature=0.7,
                max_tokens=2048,
            )
            return response.choices[0].message.content
            
        except Exception as e:
            wait_time = base_wait_time * (2 ** attempt)  # 指数退避
            print(f"智谱AI调用出错 (尝试 {attempt+1}/{max_retries}): {str(e)}")
            print(f"等待 {wait_time} 秒后重试...")
            
            if attempt < max_retries - 1:
                time.sleep(wait_time)
            else:
                print("已达到最大重试次数，返回错误信息")
                return f"调用智谱AI时遇到错误。请稍后重试。错误详情: {str(e)[:100]}..."
    
    # 这里理论上不会执行到，但是为了安全起见还是加上
    return "无法获取响应，请稍后重试。"

def search_chinanews(query, num_results=10)->list[NewsItem]:
    """
    搜索中新网并返回前n个新闻链接。对于此新闻站点的信息，
    
    参数:
        query: 搜索关键词
        num_results: 返回结果数量，默认为10
    
    返回:
        包含新闻标题和链接字典的列表
        字典参数:
        title: 新闻标题
        url: 新闻链接
    """
    url = f"https://sou.chinanews.com/search/news?q={query}"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }
    
    try:
        print(f"正在访问URL: {url}")
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        
        print(f"HTTP状态码: {response.status_code}")
        
        # 从JavaScript代码中提取docArr变量
        html_content = response.text
        doc_arr_pattern = re.compile(r'var docArr = (\[.*?\]);', re.DOTALL)
        match = doc_arr_pattern.search(html_content)
        
        if not match:
            print("未找到docArr变量")
            return []
            
        # 提取JSON字符串并解析
        doc_arr_json = match.group(1)
        try:
            doc_arr = json.loads(doc_arr_json)
            print(f"成功解析到 {len(doc_arr)} 条新闻")
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {e}")
            return []
        
        # 处理结果
        results = []
        for i, item in enumerate(doc_arr):
            if i >= num_results:
                break
                
            title = item.get("title", "")
            url = item.get("url", "")
            
            # 处理标题可能是列表的情况
            if isinstance(title, list):
                title = title[0]
                
            # 移除HTML标签
            title = re.sub(r'<[^>]+>', '', title)
            
            if url:
                results.append({"title": title, "url": url})
                print(f"提取到: {title} - {url}")
                
        return [NewsItem(title=item["title"], url=item["url"]) for item in results]
    except Exception as e:
        print(f"搜索中新网时出错: {e}")
        import traceback
        traceback.print_exc()
        return []

if __name__ == "__main__":
    # Test the LLM call
    # print("Testing LLM call...")
    # messages = [{"role": "user", "content": "简短的回答，生活的意义是什么？"}]
    # response = call_llm(messages)
    # print(f"Prompt: {messages[0]['content']}")
    # print(f"Response: {response}")

    # print("Testing ZhipuAI call...")
    # messages = [{"role": "user", "content": "简短的回答，生活的意义是什么？"}]
    # response = call_zhipuai(messages)
    # print(f"Prompt: {messages[0]['content']}")
    # print(f"Response: {response}")

    # print("Testing China News search...")
    # results = search_chinanews("科技")
    # print(f"Found {len(results)} news items:")
    # for i, result in enumerate(results, 1):
    #     print(f"{i}. {result['title']}: {result['url']}")

    pass